from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from langchain_ollama import ChatOllama
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.memory import ConversationBufferWindowMemory
from langchain_core.runnables import RunnablePassthrough
from datetime import datetime
import logging
import json
from typing import List, Dict, Optional

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="本地对话服务")


# 定义请求/响应格式
class ChatRequest(BaseModel):
    user_input: str
    session_id: str  # 用于区分不同对话会话
    max_history: Optional[int] = 5  # 可选参数，指定最大历史消息数


class Message(BaseModel):
    role: str  # "user" 或 "ai"
    content: str
    timestamp: str


class ChatResponse(BaseModel):
    ai_response: str
    conversation_id: str
    current_timestamp: str
    history: List[Message]  # 返回当前对话历史


class SessionListResponse(BaseModel):
    sessions: List[str]
    session_count: int


# 维护会话记忆和元数据（生产环境建议用Redis等分布式存储）
session_data = {
    # 结构: session_id -> {
    #     "memory": ConversationBufferWindowMemory对象,
    #     "created_at": 创建时间,
    #     "updated_at": 更新时间,
    #     "message_count": 消息数量
    # }
}


def get_memory(session_id: str, max_history: int = 5):
    """获取或创建会话记忆，带元数据管理"""
    if session_id not in session_data:
        # 创建新会话
        memory = ConversationBufferWindowMemory(
            memory_key="chat_history",
            return_messages=True,
            k=max_history
        )
        session_data[session_id] = {
            "memory": memory,
            "created_at": datetime.now().isoformat(),
            "updated_at": datetime.now().isoformat(),
            "message_count": 0
        }
        logger.info(f"创建新会话: {session_id}，最大历史消息数: {max_history}")
    else:
        # 更新现有会话的最大历史消息数
        memory = session_data[session_id]["memory"]
        if max_history != memory.k:
            memory.k = max_history
            logger.info(f"更新会话 {session_id} 的最大历史消息数为: {max_history}")

    return session_data[session_id]["memory"]


def get_session_history(session_id: str) -> List[Message]:
    """获取格式化的会话历史"""
    if session_id not in session_data:
        return []

    memory = session_data[session_id]["memory"]
    messages = memory.chat_memory.messages
    history = []

    for msg in messages:
        if msg.type == "human":
            role = "user"
        elif msg.type == "ai":
            role = "ai"
        else:
            role = "system"

        history.append(Message(
            role=role,
            content=msg.content,
            timestamp=datetime.now().isoformat()  # 实际应用中可存储每条消息的时间戳
        ))

    return history


# 初始化LLM和链
try:
    llm = ChatOllama(
        model="qwen2.5-coder:14b",
        temperature=0.7,
        timeout=30  # 超时设置
    )

    prompt = ChatPromptTemplate.from_messages([
        ("system", "你是友好的聊天助手，基于上下文进行自然、连贯的回应。如果没有上下文信息，直接回应当前问题即可。"),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{input}")
    ])

    chain = (
            RunnablePassthrough.assign(
                chat_history=lambda x: x["memory"].load_memory_variables({})["chat_history"]
            )
            | prompt
            | llm
    )
    logger.info("LLM和对话链初始化成功")
except Exception as e:
    logger.error(f"LLM初始化失败: {str(e)}")
    raise


@app.post("/chat", response_model=ChatResponse, summary="发送消息并获取回应")
def chat(request: ChatRequest):
    try:
        # 获取会话记忆
        memory = get_memory(request.session_id, request.max_history)


        # 手动更新记忆
        memory.save_context(
            {"input": request.user_input},
            {"output": response.content}
        )

        # 更新会话元数据
        if request.session_id in session_data:
            session_data[request.session_id]["updated_at"] = datetime.now().isoformat()
            session_data[request.session_id]["message_count"] += 1

        # 获取更新后的对话历史
        history = get_session_history(request.session_id)

        return ChatResponse(
            ai_response=response.content,
            conversation_id=request.session_id,
            current_timestamp=datetime.now().isoformat(),
            history=history
        )

    except Exception as e:
        logger.error(f"对话处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"对话处理错误: {str(e)}")


@app.get("/sessions", response_model=SessionListResponse, summary="获取所有会话ID")
def list_sessions():
    """返回所有活跃的会话ID列表"""
    sessions = list(session_data.keys())
    return SessionListResponse(
        sessions=sessions,
        session_count=len(sessions)
    )


@app.get("/history/{session_id}", response_model=List[Message], summary="获取指定会话的历史消息")
def get_history(session_id: str):
    """获取指定会话的完整对话历史"""
    if session_id not in session_data:
        raise HTTPException(status_code=404, detail=f"会话 {session_id} 不存在")

    return get_session_history(session_id)


@app.delete("/sessions/{session_id}", summary="删除指定会话")
def delete_session(session_id: str):
    """删除指定的会话及其历史记录"""
    if session_id not in session_data:
        raise HTTPException(status_code=404, detail=f"会话 {session_id} 不存在")

    del session_data[session_id]
    logger.info(f"删除会话: {session_id}")
    return {"message": f"会话 {session_id} 已成功删除"}


@app.get("/session-info/{session_id}", summary="获取会话信息")
def get_session_info(session_id: str):
    """获取指定会话的元数据信息"""
    if session_id not in session_data:
        raise HTTPException(status_code=404, detail=f"会话 {session_id} 不存在")

    info = session_data[session_id]
    return {
        "session_id": session_id,
        "created_at": info["created_at"],
        "updated_at": info["updated_at"],
        "message_count": info["message_count"],
        "max_history": info["memory"].k
    }